#ifndef FLUID_SOLVER_HEADER
#define FLUID_SOLVER_HEADER

#include "Field.h"
#include "glm/glm.hpp"
#include "HaarSelector.h"

template<typename T>
int DecompField( const Field<T> & inField,
	Field<T> & outWaveField)
{
	int decompDim = (inField.GetDimX()+2)*0.5f;
	int inFieldDim = inField.GetDimX();
	Field<T> tmpOutField = Field<T>(decompDim, decompDim);
	Field<T> tmpField(inField.GetDimX(), inField.GetDimX());

	if (decompDim < 10){
		return 0;
	}

	// Row decomp
	for (int r = 0; r < inField.GetDimX(); ++r){
			tmpField(r,0) =
				inField(r,0);

			tmpField(r,1) =
				-0.5f*inField(r,0) + 
				inField(r,1) + 
				(3.0f/4)*inField(r,2) - 
				(1.0f/4)*inField(r,3);

			int j = 2;
			for (int i = 4; i < inField.GetDimX()-5; i+=2){
				tmpField(r,j) = 
					(-1.0f/4.0f)*inField(r,i) +
					(3.0f/4.0f)*inField(r,i+1) +
					(3.0f/4.0f)*inField(r,i+2) -
					(1.0f/4.0f)*inField(r,i+3);
				j++;
			}
			tmpField(r,j) = -(1.0f/4.0f)*inField(r,inFieldDim-4) +
							 (3.0f/4.0f)*inField(r,inFieldDim-3) +
									     inField(r,inFieldDim-2) +
							-(1.0f/2.0f)*inField(r,inFieldDim-1);
			tmpField(r,j+1) = inField(r,inFieldDim-1);
	}

	// col
	for (int c = 0; c < decompDim; ++c){
			tmpOutField(0,c) = tmpField(0,c);

			tmpOutField(1,c) =
				-0.5f*tmpField(0,c) + 
				      tmpField(1,c) + 
				(3.0f/4)*tmpField(2,c) - 
				(1.0f/4)*tmpField(3,c);

			int j = 2;
			for (int i = 4; i < tmpField.GetDimX()-5; i+=2){
				tmpOutField(j,c) = (-1.0f/4.0f)*tmpField(i,c) +
									(3.0f/4.0f)*tmpField(i+1,c) +
									(3.0f/4.0f)*tmpField(i+2,c) -
									(1.0f/4.0f)*tmpField(i+3,c);
				j++;
			}
			tmpOutField(j,c) = -(1.0f/4.0f)*tmpField(inFieldDim-4,c) +
							    (3.0f/4.0f)*tmpField(inFieldDim-3,c) +
										    tmpField(inFieldDim-2,c) +
							   -(1.0f/2)*tmpField(inFieldDim-1,c);

			tmpOutField(j+1,c) = tmpField(inFieldDim-1,c);
	}

	outWaveField = tmpOutField;

	return 1;
}
	
template<typename T>
int DecompWavelets( const Field<T> & inField,
	Field<T> & outWaveField)
{
	int decompDim = (inField.GetDimX()-2)*0.5f;
	int inFieldDim = inField.GetDimX();
	Field<T> tmpOutField = Field<T>(decompDim, decompDim);
	Field<T> tmpField(inField.GetDimX(), inField.GetDimX());

	if (decompDim < 10){
		return 0;
	}

	// Row decomp
	for (int r = 0; r < inField.GetDimX(); ++r){
			tmpField(r,0) =
				-(1.0f/2.0f)*inField(r,0) + 
							 inField(r,1) + 
				 -(3.0f/4)*inField(r,2) + 
				 (1.0f/4)*inField(r,3);

			tmpField(r,1) =
				-(1.0f/4)*inField(r,2) + 
				(3.0f/4.0f)*inField(r,3) +
				-(3.0f/4.0f)*inField(r,4) + 
				(1.0f/4)*inField(r,5);

			int j = 2;
			for (int i = 4; i < inField.GetDimX()-5; i+=2){
				tmpField(r,j) = 
					(1.0f/4.0f)*inField(r,i) +
					-(3.0f/4.0f)*inField(r,i+1) +
					(3.0f/4.0f)*inField(r,i+2) +
					-(1.0f/4.0f)*inField(r,i+3);
				j++;
			}
			tmpField(r,j) =  (1.0f/4.0f)*inField(r,inFieldDim-4) +
							 -(3.0f/4.0f)*inField(r,inFieldDim-3) +
									     inField(r,inFieldDim-2) +
							-(1.0f/2.0f)*inField(r,inFieldDim-1);

	}

	// col
	for (int c = 0; c < decompDim; ++c){
			tmpOutField(0,c) =
				-0.5f*tmpField(0,c) + 
				      tmpField(1,c) + 
				-(3.0f/4)*tmpField(2,c) + 
				(1.0f/4)*tmpField(3,c);

			tmpOutField(1,c) =
				-(1.0f/4)*tmpField(2,c) + 
				 (3.0f/4.0f)*tmpField(3,c) +
				 -(3.0f/4.0f)*tmpField(4,c) + 
				 (1.0f/4)*tmpField(5,c);

			int j = 2;
			for (int i = 4; i < tmpField.GetDimX()-5; i+=2){
				tmpOutField(j,c) = 
					(1.0f/4.0f)*tmpField(c,i) +
					-(3.0f/4.0f)*tmpField(i+1,c) +
					(3.0f/4.0f)*tmpField(i+2,c) +
					-(1.0f/4.0f)*tmpField(i+3,c);
				j++;
			}
			tmpOutField(j,c) = (1.0f/4.0f)*tmpField(inFieldDim-4,c) +
							   -(3.0f/4.0f)*tmpField(inFieldDim-3,c) +
									       tmpField(inFieldDim-2,c) +
							  -(1.0f/2.0f)*tmpField(inFieldDim-1,c);
	}

	outWaveField = tmpOutField;

	return 1;
}
	


class FluidSolver {
public :

	enum InterpMode { BILIN, BICUBIC };

	FluidSolver(int w, int h, float dt = 0.03f);

	void SetViscosity( float visc );

	void DiffuseVelocity( );
	void AdvectVelocity( Field<glm::vec2> * inField = NULL);
	void ProjectVelocity(  );

	void AdvectField(Field<glm::vec2> * inField = NULL);

	void SetInterpMode( InterpMode );

	void SetCurrField( Field<glm::vec2> & inField ) {
		*this->GetCurrField() = inField;
	}

	void SetTimeStep( float dt );

	void SwapFields();

	void SwapScalarFields();

	void Update();

	Field<glm::vec2> * GetCurrField();

	Field<glm::vec2> * GetNewField();

	Field<float> * GetCurrScalarField();

	Field<float> * GetNewScalarField();

	// Uses second-order runge-kutta scheme
	void TraceParticle(
		const glm::vec2 & pos, 
		float deltaT, 
		const Field< glm::vec2 > & velField,
		glm::vec2 & out )
	{
		// Sample init velocity.
		glm::vec2 v0;
		velField.Sample(pos[0],pos[1], v0);

		// Update position
		glm::vec2 p1;
		p1 = pos + (v0 * deltaT);

	
		// sample velocity at new position from initial position
		glm::vec2 v1;
		velField.Sample(p1[0], p1[1], v1);

		// Take average between two positions.
		out = 0.5f*(pos + (v1*deltaT)) + 0.5f*(p1);
	}

	void SetSelector( HaarSelector * inSelector ){
		mHaarSelector = inSelector;
	}

private:
	float mVisc;
	float mDt;

	void CancleBoundary( Field<glm::vec2> & inField);
	void CancleBoundary( std::vector<float> & inVec);
	void CancleBoundary( Field<float> & inVec);

	void PoissonSolver( float * laplacA,
		unsigned int iters, 
		Field<glm::vec2> & b, 
		Field<glm::vec2> & outField );

	void PoissonSolverScalar(float * laplacA,
	unsigned int w, 
	unsigned int h,
	unsigned int iters,
	const Field<float> & b,
	Field<float> & out);


	void Clamp2f(float * vect, float min, float max) {
		for (int c = 0; c < 2; ++c){
			vect[c] = std::max(std::min(max, vect[c]), min);
		}
	}

	Field<glm::vec2> * mFieldA;
	Field<glm::vec2> * mFieldB;

	Field<float> * mScalarFieldA;
	Field<float> * mScalarFieldB;

	// XXX Testing Haar
	HaarSelector * mHaarSelector;
};

#endif